#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import json
import random
from datetime import datetime
from faker import Faker
import time
import sys
from kafka import KafkaProducer
from kafka.errors import KafkaError
from decimal import Decimal

# ------------------ 配置 ------------------
KAFKA_BOOTSTRAP_SERVERS = '192.168.177.152:9092'
KAFKA_TOPIC = 'loan-credit-events'

# 初始化 Faker
fake = Faker('zh_CN')

# App 列表
APPS = [
    "微信", "支付宝", "京东", "淘宝", "美团", "滴滴出行", "抖音", "快手",
    "高德地图", "百度地图", "银行APP", "知乎", "微博", "小红书", "拼多多"
]

def random_imei():
    return ''.join(random.choices('0123456789', k=15))

def random_mac():
    return ':'.join(['%02x' % random.randint(0, 255) for _ in range(6)])

def random_ip():
    return f"{random.randint(1, 254)}.{random.randint(0, 255)}.{random.randint(0, 255)}.{random.randint(1, 254)}"

def generate_record():
    name = fake.name()
    gender = random.choice(['男', '女'])
    age = random.randint(18, 65)
    id_card = fake.ssn(min_age=18, max_age=65)
    phone = fake.phone_number()

    return {
        "timestamp": datetime.now().isoformat(),
        "customer_name": name,
        "customer_id": f"CUST{random.randint(10000000, 99999999)}",
        "gender": gender,
        "age": age,
        "id_card": id_card,
        "nationality": "中国",
        "device_info": {
            "ip": random_ip(),
            "imei": random_imei(),
            "mac": random_mac(),
            "latitude": round(fake.latitude(), 6),
            "longitude": round(fake.longitude(), 6),
            "installed_apps": random.sample(APPS, random.randint(3, 12))
        },
        "contact_info": {
            "mobile": phone,
            "work_mobile": fake.phone_number(),
            "emergency_contact_mobile": fake.phone_number(),
            "home_address": fake.address(),
            "work_address": f"{fake.city()}高新区{fake.company()}办公楼",
            "registered_address": fake.address()
        }
    }

def custom_serializer(obj):
    if isinstance(obj, Decimal):
        return str(obj)
    elif isinstance(obj, dict):
        return {k: custom_serializer(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [custom_serializer(i) for i in obj]
    else:
        return obj

def main(total_count):
    # 创建 KafkaProducer
    try:
        producer = KafkaProducer(
            bootstrap_servers=[KAFKA_BOOTSTRAP_SERVERS],
            value_serializer=lambda v: json.dumps(custom_serializer(v), ensure_ascii=False).encode('utf-8'),
            key_serializer=lambda k: k.encode('utf-8') if k else None,
            # acks='1',  # 等待 leader 确认
            retries=3,
            linger_ms=50,  # 批量延迟
            batch_size=16384,
            client_id='loan-simulator'
        )
    except Exception as e:
        print(f"[ERROR] Failed to connect to Kafka: {e}")
        sys.exit(1)

    interval = 3600.0 / total_count  # 分布在 1 小时内

    success_count = 0
    fail_count = 0

    for i in range(total_count):
        record = generate_record()
        print(record)
        key = record['customer_id']

        try:
            # 发送消息（异步）
            future = producer.send(
                topic=KAFKA_TOPIC,
                key=key,
                value=record
            )
            # 可选：同步等待发送结果（会降低性能）
            # future.get(timeout=5)

            success_count += 1
            #print(success_count)
        except KafkaError as e:
            print(f"[ERROR] Failed to send record {key}: {e}")
            fail_count += 1
        except Exception as e:
            print(f"[ERROR] Unexpected error: {e}")
            fail_count += 1

        # 休眠（模拟持续流入）
        if i < total_count - 1:
            jitter = random.uniform(-0.5, 1.0)
            sleep_time = max(0.1, interval + jitter)
            time.sleep(sleep_time)

    # 关闭生产者
    producer.close(timeout=10)
    print(f"[INFO] Done. Sent: {success_count}, Failed: {fail_count}")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python generate_kafka.py <number_of_records>")
        sys.exit(1)
    try:
        num = int(sys.argv[1])
        main(num)
    except KeyboardInterrupt:
        print("\n[INFO] Stopped by user.")
    except Exception as e:
        print(f"[ERROR] {e}")
        sys.exit(1)
